import numpy as np
import matplotlib.pyplot as plt
from matplotlib import collections  as mc
from math import cos, sin
from pathlib import Path

figures_i = 0
figures_N = 40


def key_press_event(event):
    global figures_i
    fig = event.canvas.figure

    if event.key == 'q' or event.key == 'escape':
        plt.close(event.canvas.figure)
        return

    if event.key == 'right':
        figures_i = (figures_i + 1) % figures_N
    elif event.key == 'left':
        figures_i = (figures_i - 1) % figures_N

    fig.clear()
    my_plot(fig, figures_i)
    plt.draw()


def my_plot(fig, figures_i):
    ax = fig.add_subplot(111)

    iteration = str(figures_i).zfill(3)

    
    file_path = Path("../../output/model_landing_3dof/iteration" + iteration + "_X.txt")

    if file_path.exists():
        print('iteration',iteration,':plot')
    else:
        print('iteration',iteration,':no data')
        return

    X = np.loadtxt("../../output/model_landing_3dof/iteration" + iteration + "_X.txt")
    U = np.loadtxt("../../output/model_landing_3dof/iteration" + iteration + "_U.txt")

    ax.set_xlabel('X')
    ax.set_ylabel('Y')

    lines = []
    line_colors = []

    K = X.shape[1]

    for k in range(K):
        rx = X[0, k]
        ry = X[1, k]
        vx = X[2, k]
        vy = X[3, k]
        theta = X[4, k]
        throttle = U[0, k]
        gimbal = U[1, k]

        # speed vector
        speed_scale = 0.8
        lines.append([(rx, ry), (rx + speed_scale * vx, ry + speed_scale * vy)])
        line_colors.append((0, 1, 0, 1)) # 绿色：用于表示速度向量。

        # attitude vector
        heading_scale = 100
        c_theta = heading_scale * cos(theta)
        s_theta = heading_scale * sin(theta)
        lines.append([(rx, ry), (rx + s_theta, ry + c_theta)])
        line_colors.append((0, 0, 1, 1)) # 蓝色：用于表示姿态向量。

        # thrust vector
        throttle_scale = 80
        Tx = throttle_scale * throttle * sin(theta + gimbal)
        Ty = throttle_scale * throttle * cos(theta + gimbal)
        lines.append([(rx, ry), (rx - Tx, ry - Ty)])
        line_colors.append((1, 0, 0, 1)) # 红色：用于表示推力向量。

    lc = mc.LineCollection(lines, colors=line_colors, linewidths=1.5)

    ax.add_collection(lc)
    ax.axis('equal')
    ax.set_title("iter " + str(figures_i))


def main():
    fig = plt.figure(figsize=(10, 12))
    my_plot(fig, figures_i)
    cid = fig.canvas.mpl_connect('key_press_event', key_press_event)
    plt.show()


if __name__ == '__main__': main()
